package com.example.harmonet.harmtorch;


public class LSTM implements Layer{
    int input_size, hidden_size, num_layers;
    boolean bias;

    float[][] Wii, Whi, Wif, Whf, Wig, Whg, Wio, Who;
    float[] bii, bhi, bif, bhf, big, bhg, bio, bho;
    float[] it, ft, gt, ot, cell, h_last;

    public LSTM(int input, int hidden, int num_ls, boolean bia){
        input_size = input;
        hidden_size = hidden;
        num_layers = num_ls;
        bias = bia;

        Wii = new float[hidden_size][input_size];
        Whi = new float[hidden_size][hidden_size];
        Wif = new float[hidden_size][input_size];
        Whf = new float[hidden_size][hidden_size];
        Wig = new float[hidden_size][input_size];
        Whg = new float[hidden_size][hidden_size];
        Wio = new float[hidden_size][input_size];
        Who = new float[hidden_size][hidden_size];

        bii = new float[hidden_size];
        bhi = new float[hidden_size];
        bif = new float[hidden_size];
        bhf = new float[hidden_size];
        big = new float[hidden_size];
        bhg = new float[hidden_size];
        bio = new float[hidden_size];
        bho = new float[hidden_size];

        it = new float[hidden_size];
        ft = new float[hidden_size];
        gt = new float[hidden_size];
        ot = new float[hidden_size];
        cell = new float[hidden_size];
        h_last = new float[hidden_size];
    }

    private static float sigmoid(float val){
        return (float)1.0 / (1 + (float)Math.exp(-val));
    }

    private void doubleForInForward(Tensor in, float[][] wi, float[][] wh, float[] bi, float[] bh, float[] t){
        for(int i = 0; i < hidden_size; i++){
            t[i] = 0;
            for(int j = 0; j < input_size; j++){
                t[i] += wi[i][j]*in.tensor()[j];
            }
            for(int j = 0; j < hidden_size; j++){
                t[i] += wh[i][j]*h_last[j];
            }
            t[i] += bi[i] + bh[i];
            t[i] = sigmoid(t[i]);
        }
    }

    @Override
    public Tensor forward(Tensor in){
        Tensor out = new Tensor(hidden_size);
        float term = 0;

        /*
        for(int i = 0; i < hidden_size; i++){
            it[i] = 0;
            for(int j = 0; j < input_size; j++){
                it[i] += Wii[i][j]*in.tensor()[j];
            }
            for(int j = 0; j < hidden_size; j++){
                it[i] += Whi[i][j]*h_last[j];
            }
            it[i] += bii[i] + bhi[i];
            it[i] = sigmoid(it[i]);
        }

        for(int i = 0; i < hidden_size; i++){
            ft[i] = 0;
            for(int j = 0; j < input_size; j++){
                ft[i] += Wif[i][j]*in.tensor()[j];
            }
            for(int j = 0; j < hidden_size; j++){
                ft[i] += Whf[i][j]*h_last[j];
            }
            ft[i] += bif[i] + bhf[i];
            ft[i] = sigmoid(ft[i]);
        }

        for(int i = 0; i < hidden_size; i++){
            gt[i] = 0;
            for(int j = 0; j < input_size; j++){
                gt[i] += Wig[i][j]*in.tensor()[j];
            }
            for(int j = 0; j < hidden_size; j++){
                gt[i] += Whg[i][j]*h_last[j];
            }
            gt[i] += big[i] + bhg[i];
            gt[i] = (float)Math.tanh(gt[i]);
        }

        for(int i = 0; i < hidden_size; i++){
            ot[i] = 0;
            for(int j = 0; j < input_size; j++){
                ot[i] += Wio[i][j]*in.tensor()[j];
            }
            for(int j = 0; j < hidden_size; j++){
                ot[i] += Who[i][j]*h_last[j];
            }
            ot[i] += bio[i] + bho[i];
            ot[i] = sigmoid(ot[i]);
        }
        */

        doubleForInForward(in, Wii, Whi, bii, bhi, it);
        doubleForInForward(in, Wif, Whf, bif, bhf, ft);
        doubleForInForward(in, Wig, Whg, big, bhg, gt);
        doubleForInForward(in, Wio, Who, bio, bho, ot);

        for(int i = 0; i < hidden_size; i++){
            cell[i] = ft[i]*cell[i] + it[i]*gt[i];
            out.tensor()[i] = ot[i]*(float)Math.tanh(cell[i]);
        }

        h_last = out.tensor().clone();
        return out;
    }

    @Override
    public int getParam(){
        return 4*hidden_size*input_size + 4*hidden_size*hidden_size + 10*hidden_size;
    }

    private int doubleForAssign(int pos, float[][] arr, float[] par, int size1, int size2){
        for(int i = 0; i < size1; i++)
            for(int j = 0; j < size2; j++){
                arr[i][j] = par[pos];
                pos++;
            }
        return pos;
    }

    private int singleForAssign(int pos, float[] arr, float[] par, int size1){
        for(int i = 0; i < size1; i++){
            arr[i] = par[pos];
            pos++;
        }
        return pos;
    }

    @Override
    public void init(float[] par){
        // W, b, cell0, h0. length of par: 4*hidden*input + 4*hidden^2 + 10*hidden
        int curPos = 0;

        // information gate (input gate)
        /*
        for(int i = 0; i < hidden_size; i++){
            for(int j = 0; j < input_size; j++){
                Wii[i][j] = par[curPos++];
            }
        }

        for(int i = 0; i < hidden_size; i++){
            for(int j = 0; j < hidden_size; j++){
                Whi[i][j] = par[curPos++];
            }
        }
        */
        curPos = doubleForAssign(curPos, Wii, par, hidden_size, input_size);
        curPos = doubleForAssign(curPos, Whi, par, hidden_size, hidden_size);

        // forget gate (f)
        /*
        for(int i = 0; i < hidden_size; i++){
            for(int j = 0; j < input_size; j++){
                Wif[i][j] = par[curPos++];
            }
        }

        for(int i = 0; i < hidden_size; i++){
            for(int j = 0; j < hidden_size; j++){
                Whf[i][j] = par[curPos++];
            }
        }
        */
        curPos = doubleForAssign(curPos, Wif, par, hidden_size, input_size);
        curPos = doubleForAssign(curPos, Whf, par, hidden_size, hidden_size);

        // intrinsic weight matrix (g)
        /*
        for(int i = 0; i < hidden_size; i++){
            for(int j = 0; j < input_size; j++){
                Wig[i][j] = par[curPos++];
            }
        }

        for(int i = 0; i < hidden_size; i++){
            for(int j = 0; j < hidden_size; j++){
                Whg[i][j] = par[curPos++];
            }
        }
        */

        curPos = doubleForAssign(curPos, Wig, par, hidden_size, input_size);
        curPos = doubleForAssign(curPos, Whg, par, hidden_size, hidden_size);


        // out gate (o)
        /*
        for(int i = 0; i < hidden_size; i++){
            for(int j = 0; j < input_size; j++){
                Wio[i][j] = par[curPos++];
            }
        }

        for(int i = 0; i < hidden_size; i++){
            for(int j = 0; j < hidden_size; j++){
                Who[i][j] = par[curPos++];
            }
        }
        */
        curPos = doubleForAssign(curPos, Wif, par, hidden_size, input_size);
        curPos = doubleForAssign(curPos, Whf, par, hidden_size, hidden_size);


        // get b
        /*
        for(int i = 0; i < hidden_size; i++){
            bii[i] = par[curPos++];
        }
        for(int i = 0; i < hidden_size; i++){
            bhi[i] = par[curPos++];
        }
        for(int i = 0; i < hidden_size; i++){
            bif[i] = par[curPos++];
        }
        for(int i = 0; i < hidden_size; i++){
            bhf[i] = par[curPos++];
        }
        for(int i = 0; i < hidden_size; i++){
            big[i] = par[curPos++];
        }
        for(int i = 0; i < hidden_size; i++){
            bhg[i] = par[curPos++];
        }
        for(int i = 0; i < hidden_size; i++){
            bio[i] = par[curPos++];
        }
        for(int i = 0; i < hidden_size; i++){
            bho[i] = par[curPos++];
        }
        */

        curPos = singleForAssign(curPos, bii, par, hidden_size);
        curPos = singleForAssign(curPos, bhi, par, hidden_size);
        curPos = singleForAssign(curPos, bif, par, hidden_size);
        curPos = singleForAssign(curPos, bhf, par, hidden_size);
        curPos = singleForAssign(curPos, big, par, hidden_size);
        curPos = singleForAssign(curPos, bhg, par, hidden_size);
        curPos = singleForAssign(curPos, bio, par, hidden_size);
        curPos = singleForAssign(curPos, bho, par, hidden_size);


        /*
        // cell @ t = 0
        for(int i = 0; i < hidden_size; i++){
            cell[i] = par[curPos++];
        }

        // h @ t = 0
        for(int i = 0; i < hidden_size; i++){
            h_last[i] = par[curPos++];
        }
        */
        curPos = singleForAssign(curPos, cell, par, hidden_size);
        curPos = singleForAssign(curPos, h_last, par, hidden_size);

    }
}
